# src/modeling.py

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, classification_report
from xgboost import XGBClassifier
import joblib
import os

def train_retention_model(user_df, label_col='label', save_path='results/xgb_model.pkl'):
    # 筛选特征和标签
    feature_cols = [col for col in user_df.columns if col not in ['user_id', 'label']]
    X = user_df[feature_cols]
    y = user_df[label_col]

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    # 模型训练
    model = XGBClassifier(use_label_encoder=False, eval_metric='logloss')
    model.fit(X_train, y_train)

    # 模型评估
    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)

    print(f"✅ 模型训练完成！准确率: {acc:.4f}, F1-score: {f1:.4f}")
    print("📋 详细评估报告：\n", classification_report(y_test, y_pred))

    # 模型保存
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    joblib.dump(model, save_path)
    print(f"🧠 模型已保存至 {save_path}")
